# python-resmon - Python tools for Resmon monitoring and the Circonus API
# Copyright (c) 2010 Crown coypright
# 
# This file is part of python-resmon.
# 
# python-resmon is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# python-resmon is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with python-resmon.  If not, see <http://www.gnu.org/licenses/>.

"""Implements a webservice that can be queried for the current state of the infrastructure"""

import time

import threading
from Queue import Queue, Empty

from alert import *
from checker import ResmonBroker, ResmonChecker, CheckThread 
from ruleparse.loader import load


class StatusWorker(CheckThread):
    """A worker that takes check jobs from a queue and enqueues the results.

    Based on CheckThread but only shares its setup; it differs in two main ways:
      - it doesn't die if there are no jobs, it just waits 2s
      - it enqueues lists of results, not individual alerts.
    """
    def run(self):
        while self.keeprunning:
            host = self.hostqueue.get(block=True)
#            print 'C:', host.hostname
            alerts = list(self.checker.check_host(host))
            self.outqueue.put((host, alerts))


class TaskQueue(Queue):
    """A queue where tasks can only be enqueued once
    until they are marked as completed.

    Tasks must be hashable.
    """
    class DuplicateTask(Exception):
        """The task is already queued."""

    def __init__(self, maxsize=None):
        Queue.__init__(self, maxsize)
        self.lock = threading.Lock()
        self.queued_tasks = set()

    def put(self, task, block=True, timeout=None):
        """Put the task into the queue.
        
        If the task is already in the queue, raises TaskQueue.DuplicateTask.
        """
        with self.lock:
            if task in self.queued_tasks:
                raise TaskQueue.DuplicateTask(task)
            self.queued_tasks.add(task)
            Queue.put(self, task, block, timeout)

    def done(self, task):
        """Mark the task as done, so that it can be requeued."""
        with self.lock:
            self.queued_tasks.remove(task)



class StatusService(threading.Thread):
    """Keep running checks so as to maintain an up-to-date picture of the state of the infrastructure.
    """
    WORKERS = 8
    CHECK_INTERVAL = 60
    RETRY_INTERVAL = 6

    def __init__(self, rulefile='rules.cm', broker=ResmonBroker()):
        super(StatusService, self).__init__()

        self.checker = ResmonChecker(rulefile, broker)
        self.daemon = True

        self.alerts = {} # alerts by host
        self.attempts = {} # retry count on failed checks, by host

        self.lock = threading.Lock()

        self.outqueue = Queue() # queue of events to process back in
        self.taskqueue = TaskQueue(maxsize=1000) # queue of jobs to delegate to workers

        self.keeprunning = True
        self.pool = []

        # Map of hostnames to the next time to check that host
        self.schedule = {}

    def generate_schedule(self):
        """Generate a schedule by spreading the checks evenly over CHECK_INTERVAL"""
        hosts = self.checker.model
        interval = self.CHECK_INTERVAL / float(len(hosts))
        for i, host in enumerate(hosts):
            self.schedule_check(host, i * interval)

    def spawn_workers(self):
        if self.pool:
            return
        self.pool = [StatusWorker(self.checker, self.taskqueue, self.outqueue) for i in range(self.WORKERS)]

    def queue_checks(self):
        """Enqueue checks that are either not scheduled or whose scheduled time
        has past.
        """
        now = time.time()
        for host in self.checker.model:
            if self.schedule.get(host.hostname, 0) < now:
                try:
                    self.taskqueue.put(host)
                except TaskQueue.DuplicateTask:
                    pass
#                else:
#                    print 'Q:', host.hostname

    def get_alerts(self):
        """Return a list of all current alerts.
        
        Blocks only as long as necessary to obtain a lock on the alert table.
        """
        alerts = []
        with self.lock:
            alert_lists = self.alerts.values()
        for host_alerts in alert_lists:
            alerts += host_alerts
        return alerts

    def should_retry(self, host, alerts):
        if len(alerts) == 1:
            alert = alerts[0]
            if isinstance(alert, ResmonErrorAlert):
                return True
        return False

    def schedule_check(self, host, delay):
        """Schedule a check for host, delay seconds in the future"""
        self.schedule[host.hostname] = time.time() + delay

    def collect_results(self):
        """Records the results from the worker thread in the dictionary.
        """
        count = 0
        while True:
            try:
                host, alerts = self.outqueue.get_nowait()
            except Empty:
                return count

            self.taskqueue.done(host)
            count += 1

            if alerts and self.should_retry(host, alerts):
                attempts = self.attempts.get(host.hostname, 0) + 1
                self.attempts[host.hostname] = attempts 
                print "Failed attempts for", host.hostname, "=", attempts
                # Only update the machine status if we have had three or more failed attempts
                if attempts >= 3:
                    with self.lock:
                        self.alerts[host] = alerts

                # Increasing retry time
                self.schedule_check(host, min(self.RETRY_INTERVAL * attempts, 60))
            else:
                if self.attempts.get(host.hostname, 0):
                    print host.hostname, "succeeded after", self.attempts[host.hostname], "attempts"
                    self.attempts[host.hostname] = 0
                with self.lock:
                    self.alerts[host] = alerts

                self.schedule_check(host, self.CHECK_INTERVAL)

            # TODO: possibly notify an event dispatcher if the alerts have changed 

    def run(self):
        self.keeprunning = True
        self.spawn_workers()
        try:
            while self.keeprunning:
                self.queue_checks()
                num_results = self.collect_results()
#                print "R:", num_results
                if not num_results:
                    time.sleep(2)
        except (KeyboardInterrupt, SystemExit):
            for worker in self.pool:
                worker.stop()
            raise
